import argparse
import multiprocessing
import os
import pickle
import sys
import warnings
import xml.etree.ElementTree as ET
from copy import deepcopy
from multiprocessing import Manager, Pipe
from operator import truth
from threading import Thread

import cv2
import matplotlib.pyplot as plt
import numpy as np
import scipy.io as sio
import spectral.io.envi as envi
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from matplotlib.backends.backend_pdf import PdfPages
from numpy import flip
from spectral import imshow, view_cube
from sshkeyboard import listen_keyboard

from AutoGPU import autoGPU
from GAN_training_utils import DataResult, TrainProcess, setup_seed
# from AutoGPU import autoGPU
from models import (_1DCNN, _2DCNN, _3DCNN, _3DCNN_1DCNN, _3DCNN_AM, PURE2DCNN,
                    PURE3DCNN, PURE3DCNN_2AM, SAE, SAE_AM, DBDA_network,
                    HamidaEtAl, LeeEtAl, SSRN_network, _2dCNN, myknn, mysvm)
from myTrans2 import Generator
from NViT import ViT as NViT
from utils import DataPreProcess, myplot, plot, setpath, splitdata

if __name__ == '__main__':
    while True:
        try:
            autoGPU(1, 5000)
            break
        except Exception as e :
            print(e)
            pass

    gen_model = Generator().to('cuda')

        # dis_model.load_state_dict(torch.load(resultpath + '1bestmodel.pth'))
    t =  torch.load('pathology/032370b/roi2/Split/proportion/Tr_0.01/Va_0.01/Te_0.98/1/result/TransGan/genmodel.pth')
    gen_model.load_state_dict({k.replace('module.',''):v for k,v in t.items()}, strict=True)

    sample = np.zeros((0, 60, 9, 9))
    for i in range(100): 
    # processeddata['train'].patch.to('cuda')
        noise = torch.FloatTensor(1, 50).to('cuda')
        noise.data.resize_(1, 50).normal_(0, 1)
                    
        class_onehot = torch.ones((1, 10)).to('cuda')
                
        noise[np.arange(1), :10] = (class_onehot * 0).unsqueeze(1)
        out = gen_model(noise)
        out = out.detach().cpu().numpy()
        sample = np.vstack((sample, out))

    for i in sample[:,:,4,4].sum(axis=0):
        print(i)

    # with PdfPages('fake样本.pdf') as pdf:
    #     fig = plt.figure()
    #     plt.imshow(out_color)
    #     height, width, channels = out_color.shape
    #     fig.set_size_inches(width / 100.0, height / 100.0)
    #     plt.gca().xaxis.set_major_locator(plt.NullLocator())
    #     plt.gca().yaxis.set_major_locator(plt.NullLocator())
    #     plt.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0, wspace=0)
    #     plt.margins(0, 0)
    #     plt.axis('off')
    #     plt.xticks([])
    #     plt.yticks([])  
    #     pdf.savefig(bbox_inches = 'tight')  # saves the current figure into a pdf page
    #     plt.close()
   
    # with open(resultpath + 'result.pkl', 'wb') as f:
    #     pickle.dump(T.test_result, f, pickle.HIGHEST_PROTOCOL)
    # myplot(processeddata, IMAGE, imagepath, T.test_result)
    print('end')
